"""This module provides functionalities to work with image's color.
- color_stretch_linear
- color_stretch_normal
- color_stretch_percentage
- histogram_equalization
- histogram_matching

Note: Many RS softwares also have stretch via gamma/log/sigmoid implemented. Will implement them here as requested.
"""

from typing import Optional, List, Union, Tuple
import numpy as np
import numpy.typing as npt

from skimage import exposure

from common.typing_utils import numeric
from .img_shape import BandOrderType


def color_stretch_linear(
    nda_img: np.ndarray,
    src_min: Optional[Union[List[numeric], numeric]] = None,
    src_max: Optional[Union[List[numeric], numeric]] = None,
    dest_min: Optional[Union[List[numeric], numeric]] = None,
    dest_max: Optional[Union[List[numeric], numeric]] = None,
    band_order_type: BandOrderType = BandOrderType.GDAL,
    output_dtype: npt.DTypeLike = np.uint8,
    valid_mask: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Linearly stretch numpy array (2D single band image or 3D multi-band image)

    Parameters
    ----------
    nda_img: np.ndarray
        A numpy array (2D/3D).
    src_min: list of number, or number, or none
        minimum value of data range for stretching.
        default value: None --> the min value of the input `nda_img` is to be used.
        If a single numeric value is provided, then that value is going to be used for all bands.
    src_max list of number, or number
        maximum value of data range for stretching.
        default value: None --> the max value of the input `nda_img` is to be used.
        If a single numeric value is provided, then that value is going to be used for all bands.
    dest_min: list of number, or number
        minimum value of data range for the stretched output image.
        default value: None --> the min value of the output_dtype is to be used.
            If output_dtype is float, min value = 0.
        If a single numeric value is provided, then that value is going to be used for all bands.
    dest_max: list of number, or number
        maximum value of data range for the stretched output image.
        default value: None --> the max value of the output_dtype is to be used.
            If output_dtype is float, max value = 1.
        If a single numeric value is provided, then that value is going to be used for all bands.
    band_order_type: BandOrderType class object
        This variable defines the structure of the nda_img numpy array.
        default value: BandOrderType.GDAL --> the length of 1st dimension is equivalent to numbers of bands.
        alternative value: BandOrderType.IMAGE --> the length of 3rd dimension is equivalent to numbers of bands.
    output_dtype: numpy array dtype
        Defines data type of output numpy array
        default: np.uint8
    valid_mask : np.ndarray, optional
        A boolean numpy array with same shape as `nda_img`, True indicating the value at that place is valid.
        Notice `valid_mask` is only used for deriving `src_min` and `src_max` if they are not provided.

    Returns
    -------
    nda_img_out: np.ndarray

    Raises
    ------
    ValueError
        - dimension of lists (e.g. src_min, src_max, dest_min, dest_max) are expected to have the same length as band.

    Notes
    -----
    If `dest_min`, `dest_max` provided,
    it is the caller's duty to make sure those values fit within the supported value range of the `output_dtype`.
    It is also the caller's duty to ensure output_dtype is accurately defined.
    (e.g. With default setting, an image with data ranging between 0-1 would
    automatically be converted to uint8 (0-255) in this process.
    As a result, the values would stretches into int 0-255, instead of float 0-1.)
    """
    # define band location
    band_loc, length_width_loc = _define_band_location(band_order_type)
    # find number of bands
    if len(nda_img.shape) == 2:
        bands_num = 1
    else:
        bands_num = nda_img.shape[band_loc]
    # if input image is 2D, increase dimension to 3D, based on output_dtype setting.
    if len(nda_img.shape) == 2:
        reshape_later = True  # raise dimension now, but change it back later
        nda_img = np.expand_dims(nda_img, axis=band_loc)
    else:
        reshape_later = False
    # if stretch parameters are list, check if length matches band numbers.
    stretch_parameters = ["src_min", "src_max", "dest_min", "dest_max"]
    for check_var in stretch_parameters:
        if isinstance(eval(check_var), list):
            if len(eval(check_var)) != bands_num:
                raise ValueError(
                    (
                        f"length of {check_var} ({len(eval(check_var))})",
                        f"does not match the number of bands ({bands_num})",
                    )
                )
    # if stretch parameter are defaults, define stretch parameters according to nda_img and output_dtype
    src_min, src_max, dest_min, dest_max = _check_stretch_parameters(
        nda_img,
        length_width_loc,
        src_min,
        src_max,
        dest_min,
        dest_max,
        output_dtype,
        valid_mask,
    )
    src_min_temp, src_max_temp, dest_min_temp, dest_max_temp = None, None, None, None
    # Copy data
    nda_img_out = np.empty_like(nda_img, dtype=output_dtype)
    # loop for band
    for band_i in range(bands_num):
        if band_order_type == "gdal":
            nda_img_slice = nda_img[band_i, :, :]
        elif band_order_type == "image":
            nda_img_slice = nda_img[:, :, band_i]
        # Defines temporary variables
        (
            src_min_temp,
            src_max_temp,
            dest_min_temp,
            dest_max_temp,
        ) = _slice_stretch_parameters(band_i, src_min, src_max, dest_min, dest_max)
        nda_img_slice = exposure.rescale_intensity(
            nda_img_slice,
            in_range=(src_min_temp, src_max_temp),
            out_range=(dest_min_temp, dest_max_temp),
        )
        if band_order_type == "gdal":
            nda_img_out[band_i, :, :] = nda_img_slice
        elif band_order_type == "image":
            nda_img_out[:, :, band_i] = nda_img_slice
    # Change 3D nda_img to 2D, if nda_img were 2D to begin with
    if reshape_later:
        if band_order_type == "gdal":
            nda_img_out = nda_img_out[0, :, :]
        elif band_order_type == "image":
            nda_img_out = nda_img_out[:, :, 0]
    return nda_img_out


def color_stretch_normal(
    nda_img: np.ndarray,
    sigma: numeric = 2,
    dest_min: Optional[Union[List[numeric], numeric]] = None,
    dest_max: Optional[Union[List[numeric], numeric]] = None,
    band_order_type: BandOrderType = BandOrderType.GDAL,
    output_dtype: npt.DTypeLike = np.uint8,
    valid_mask: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Linearly stretch numpy array (2D single band image or 3D multi-band image)
    The stretch min and max range from each band is defined by X standard deviation away from mean of the band image.

    Parameters
    ----------
    nda_img: np.ndarray
        A numpy array (2D/3D).
    sigma: number
        Define number of standard deviation.
    dest_min: list of number, or number
        minimum value of data range for the stretched output image.
        default value: None --> the min value of the output_dtype is to be used.
            If output_dtype is float, min value = 0.
        If a single numeric value is provided, then that value is going to be used for all bands.
    dest_max: list of number, or number
        maximum value of data range for the stretched output image.
        default value: None --> the max value of the output_dtype is to be used.
            If output_dtype is float, max value = 1.
        If a single numeric value is provided, then that value is going to be used for all bands.
    band_order_type: BandOrderType class object
        This variable defines the structure of the nda_img numpy array.
        default value: BandOrderType.GDAL --> the length of 1st dimension is equivalent to numbers of bands.
        alternative value: BandOrderType.IMAGE --> the length of 3rd dimension is equivalent to numbers of bands.
    output_dtype: numpy array dtype
        Defines data type of output numpy array
        default: np.uint8
    valid_mask : np.ndarray, optional
        A boolean numpy array with same shape as `nda_img`, True indicating the value at that place is valid.

    Returns
    -------
    nda_img_out: np.ndarray

    Notes
    -----
    The function will automatically stretch figure by the defined range (mean +- sigma std)
    to occupy the full value range of given data type.
    It also ignore any nan value during the computation.
    """
    # define band location
    band_loc, _ = _define_band_location(band_order_type)
    # find number of bands
    if len(nda_img.shape) == 2:
        bands_num = 1
    else:
        bands_num = nda_img.shape[band_loc]
    # if input image is 2D or number of band == 1.
    if bands_num == 1:
        if valid_mask is None:
            mean_val, std_val = np.nanmean(nda_img), np.nanstd(nda_img)
        else:
            nda_img_tmp = nda_img[valid_mask]
            mean_val, std_val = np.nanmean(nda_img_tmp), np.nanstd(nda_img_tmp)
        # define src_min, src_max a single value, using percentile
        src_min, src_max = mean_val - (sigma * std_val), mean_val + (sigma * std_val)
    else:
        # define src_min, src_max as list of values, using percentile
        src_min, src_max = [], []
        for band_i in range(bands_num):
            if band_order_type == "gdal":
                if valid_mask is None:
                    nda_img_slice = nda_img[band_i, :, :]
                else:
                    nda_img_slice = nda_img[band_i, valid_mask]
            elif band_order_type == "image":
                if valid_mask is None:
                    nda_img_slice = nda_img[:, :, band_i]
                else:
                    nda_img_slice = nda_img[valid_mask, band_i]
            mean_val, std_val = np.nanmean(nda_img_slice), np.nanstd(nda_img_slice)
            # define src_min, src_max a single value, using percentile
            src_min_slice, src_max_slice = (
                mean_val - (sigma * std_val),
                mean_val + (sigma * std_val),
            )
            src_min.append(src_min_slice)
            src_max.append(src_max_slice)
    nda_img_out = color_stretch_linear(
        nda_img, src_min, src_max, dest_min, dest_max, band_order_type, output_dtype
    )
    return nda_img_out


def color_stretch_percentage(
    nda_img: np.ndarray,
    percentage: numeric = 2,
    dest_min: Optional[Union[List[numeric], numeric]] = None,
    dest_max: Optional[Union[List[numeric], numeric]] = None,
    band_order_type: BandOrderType = BandOrderType.GDAL,
    output_dtype: npt.DTypeLike = np.uint8,
    valid_mask: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Linearly stretch numpy array (2D single band image or 3D multi-band image)
    The stretch min and max range from each band is defined by percentile of the band image.

    Parameters
    ----------
    nda_img: np.ndarray
        A numpy array (2D/3D).
    percentage: number
        Define percentile for stretching.
        By default, uses the 2nd and 98th percentile to do linear stretching
    dest_min: list of number, or number
        minimum value of data range for the stretched output image.
        default value: None --> the min value of the output_dtype is to be used.
            If output_dtype is float, min value = 0.
        If a single numeric value is provided, then that value is going to be used for all bands.
    dest_max: list of number, or number
        maximum value of data range for the stretched output image.
        default value: None --> the max value of the output_dtype is to be used.
            If output_dtype is float, max value = 1.
        If a single numeric value is provided, then that value is going to be used for all bands.
    band_order_type: BandOrderType class object
        This variable defines the structure of the nda_img numpy array.
        default value: BandOrderType.GDAL --> the length of 1st dimension is equivalent to numbers of bands.
        alternative value: BandOrderType.IMAGE --> the length of 3rd dimension is equivalent to numbers of bands.
    output_dtype: numpy array dtype
        Defines data type of output numpy array
        default: np.uint8
    valid_mask : np.ndarray, optional
        A boolean numpy array with same shape as `nda_img`, True indicating the value at that place is valid.

    Returns
    -------
    nda_img_out: np.ndarray

    Raises
    ------
    ValueError
        - percentile exceed or equal to 50.

    Notes
    -----
    The function will automatically stretch figure by percentile to occupy the full value range of given data type.
    """
    # define band location
    band_loc, _ = _define_band_location(band_order_type)
    # find number of bands
    if len(nda_img.shape) == 2:
        bands_num = 1
    else:
        bands_num = nda_img.shape[band_loc]
    if percentage < 50:
        pv_1, pv_2 = percentage, 100 - percentage
    else:
        raise ValueError("Percentage value should be less than 50.")
    # if input image is 2D or number of band == 1.
    if bands_num == 1:
        # define src_min, src_max a single value, using percentile
        if valid_mask is None:
            src_min, src_max = np.percentile(nda_img, pv_1), np.percentile(
                nda_img, pv_2
            )
        else:
            nda_img_tmp = nda_img[valid_mask]
            src_min = np.percentile(nda_img_tmp, pv_1)
            src_max = np.percentile(nda_img_tmp, pv_2)
    else:
        # define src_min, src_max as list of values, using percentile
        src_min, src_max = [], []
        for band_i in range(bands_num):
            if band_order_type == "gdal":
                if valid_mask is None:
                    nda_img_slice = nda_img[band_i, :, :]
                else:
                    nda_img_slice = nda_img[band_i, valid_mask]
            elif band_order_type == "image":
                if valid_mask is None:
                    nda_img_slice = nda_img[:, :, band_i]
                else:
                    nda_img_slice = nda_img[valid_mask, band_i]
            src_min_slice, src_max_slice = (
                np.percentile(nda_img_slice, pv_1),
                np.percentile(nda_img_slice, pv_2),
            )
            src_min.append(src_min_slice)
            src_max.append(src_max_slice)
    nda_img_out = color_stretch_linear(
        nda_img, src_min, src_max, dest_min, dest_max, band_order_type, output_dtype
    )
    return nda_img_out


def histogram_equalization(
    nda_img: np.ndarray,
    nbins: int = 256,
    band_order_type: BandOrderType = BandOrderType.GDAL,
    mask: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Enhances image using the Histogram Equalization method,
    which “spreads out the most frequent intensity values” in an image.
    The equalized image has a roughly linear cumulative distribution function.

    Parameters
    ----------
    nda_img: np.ndarray
        A numpy array (2D/3D).
    nbins: int
        number of bins for histogram computation
    band_order_type: BandOrderType class object
        This variable defines the structure of the nda_img numpy array.
        default value: BandOrderType.GDAL --> the length of 1st dimension is equivalent to numbers of bands.
        alternative value: BandOrderType.IMAGE --> the length of 3rd dimension is equivalent to numbers of bands.
    mask: np.ndarray of booleans or 0s and 1s
        Same size as nda_image. Only points at which mask == True or 1 are used for equalization,
        which is applied to the whole image.

    Returns
    -------
    nda_img_out: np.ndarray

    Raises
    ------
    ValueError
        - size of nda_img not equal to size of mask

    Notes
    -----
    The output of scikit's equalize hist is in float64, used color_stretch_linear for rescaling.
    This function do not convert image to other dtype; could potentially implement it here with color_stretch_linear.
    If this function yields unnatural looking images, consider using other image stretch methods,
    There is also an adaptive histogram equalization, but have not been implemented yet.

    Reference
    --------
    Scikit-image Histogram Equalization:
    https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_equalize.html
    """
    # check size of nda_image and mask
    if (mask is not None) and (np.shape(nda_img) != np.shape(mask)):
        raise ValueError("Size of mask and nda_img do not match")

    # define band location
    band_loc, length_width_loc = _define_band_location(band_order_type)
    # find number of bands
    if len(nda_img.shape) == 2:
        bands_num = 1
    else:
        bands_num = nda_img.shape[band_loc]
    output_dtype = nda_img.dtype
    # if input image is 2D or number of band == 1.
    if bands_num == 1:
        nda_img = exposure.equalize_hist(nda_img, nbins, mask)
    else:
        nda_img_copy = np.empty_like(nda_img, dtype=np.float64)
        for band_i in range(bands_num):
            if band_order_type == "gdal":
                mask_slice = None if (mask is None) else mask[band_i, :, :]
                nda_img_copy[band_i, :, :] = exposure.equalize_hist(
                    nda_img[band_i, :, :], nbins, mask_slice
                )
            elif band_order_type == "image":
                mask_slice = None if (mask is None) else mask[:, :, band_i]
                nda_img_copy[:, :, band_i] = exposure.equalize_hist(
                    nda_img[:, :, band_i], nbins, mask_slice
                )
        nda_img = nda_img_copy
    # Rescale image back to input data type
    if output_dtype != np.float64:
        src_min, src_max, dest_min, dest_max = 0, 1, None, None
        nda_img_out = color_stretch_linear(
            nda_img, src_min, src_max, dest_min, dest_max, band_order_type, output_dtype
        )
    else:
        nda_img_out = nda_img
    return nda_img_out


def histogram_matching(
    nda_img: np.ndarray,
    nda_ref: np.ndarray,
    band_order_type: BandOrderType = BandOrderType.GDAL,
) -> np.ndarray:
    """This manipulates the pixels of an input image
    so that its histogram matches the histogram of the reference image.
    If the images have multiple channels, the matching is done independently
    for each channel, as long as the number of channels is equal in the input
    image and the reference. (Copied scikit-image example)

    Parameters
    ----------
    nda_img: np.ndarray
        A numpy array (2D/3D).
    nda_ref: np.ndarray
        A numpy array (2D/3D). Must have the same number of dimension (not shape) as nda_img
    band_order_type: BandOrderType class object
        This variable defines the structure of the nda_img numpy array.
        default value: BandOrderType.GDAL --> the length of 1st dimension is equivalent to numbers of bands.
        alternative value: BandOrderType.IMAGE --> the length of 3rd dimension is equivalent to numbers of bands.

    Returns
    -------
    nda_img_out: np.ndarray

    Raises
    ------
    ValueError
        - dimension of nda_img not equal to dimension of nda_ref
        - data type of nda_img and nda_ref must match (same value range)

    Notes
    -----
    [Future] Should add mask functionality, so histogram cam be matched between
    overlapping parts of the images.
    multichannel will be deprecated in after Scikit-image 0.19.0 update

    Reference
    --------
    Scikit-image Histogram Matching:
    https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_histogram_matching.html
    """
    # Check image data type
    if nda_img.dtype != nda_ref.dtype:
        raise ValueError(
            f"data type of nda_img ({nda_img.dtype}) and nda_ref ({nda_ref.dtype}) do not match"
        )
    # define band location
    band_loc, length_width_loc = _define_band_location(band_order_type)
    # find number of bands
    if len(nda_img.shape) == 2:
        nda_img = np.expand_dims(nda_img, axis=band_loc)
        reshape_later = True
    else:
        reshape_later = False
    if len(nda_ref.shape) == 2:
        nda_ref = np.expand_dims(nda_ref, axis=band_loc)
    # Check image shape, position here after expanding 2D inputs
    if len(nda_img.shape) != len(nda_ref.shape):
        raise ValueError(
            f"dimension of nda_img ({len(nda_img.shape)}) and nda_ref ({len(nda_ref.shape)}) do not match"
        )
    # histogram match
    nda_img_out = exposure.match_histograms(nda_img, nda_ref, multichannel=True)
    # Turn image back to 2D if necessary
    if reshape_later:
        if band_order_type == "gdal":
            nda_img_out = nda_img_out[0, :, :]
        elif band_order_type == "image":
            nda_img_out = nda_img_out[:, :, 0]
    return nda_img_out


def _define_band_location(band_order_type: BandOrderType = BandOrderType.GDAL):
    """Define band location, given image band_order_type"""
    # define band location
    if band_order_type == "gdal":
        band_loc = 0
        length_width_loc = (1, 2)
    elif band_order_type == "image":
        band_loc = 2
        length_width_loc = (0, 1)
    return (band_loc, length_width_loc)


def _check_stretch_parameters(
    nda_img: np.ndarray,
    length_width_loc: Tuple[int, int],
    src_min: Optional[Union[List[numeric], numeric]] = None,
    src_max: Optional[Union[List[numeric], numeric]] = None,
    dest_min: Optional[Union[List[numeric], numeric]] = None,
    dest_max: Optional[Union[List[numeric], numeric]] = None,
    output_dtype: npt.DTypeLike = np.uint8,
    valid_mask: Optional[np.ndarray] = None,
):
    """Check whether color stretch parameters are None,
    If None, assign value given output_dtype"""
    if valid_mask is not None:
        if src_min is None or src_max is None:
            if length_width_loc == (1, 2):
                nda_tmp = nda_img[:, valid_mask]
            else:  # length_width_loc == (0, 1)
                nda_tmp = np.moveaxis(nda_img[valid_mask, :], -1, 0)

            if src_min is None:
                src_min = list(np.min(nda_tmp, axis=1))
            if src_max is None:
                src_max = list(np.max(nda_tmp, axis=1))
    else:
        if src_min is None:
            src_min = np.amin(nda_img, axis=length_width_loc)
        if src_max is None:
            src_max = np.amax(nda_img, axis=length_width_loc)

    if dest_min is None:
        if np.issubclass_(output_dtype, np.floating):
            dest_min = 0
        else:
            dest_min = np.iinfo(output_dtype).min
    if dest_max is None:
        if np.issubclass_(output_dtype, np.floating):
            dest_max = 1
        else:
            dest_max = np.iinfo(output_dtype).max
    return (src_min, src_max, dest_min, dest_max)


def _slice_stretch_parameters(
    band_i: int = 0,
    src_min: Optional[Union[List[numeric], numeric]] = None,
    src_max: Optional[Union[List[numeric], numeric]] = None,
    dest_min: Optional[Union[List[numeric], numeric]] = None,
    dest_max: Optional[Union[List[numeric], numeric]] = None,
) -> Tuple[Optional[numeric], Optional[numeric], Optional[numeric], Optional[numeric]]:
    """Define temporary stretching parameter"""
    if isinstance(src_min, list) or isinstance(src_min, np.ndarray):
        src_min_temp = src_min[band_i]
    else:
        src_min_temp = src_min

    if isinstance(src_max, list) or isinstance(src_max, np.ndarray):
        src_max_temp = src_max[band_i]
    else:
        src_max_temp = src_max

    if isinstance(dest_min, list) or isinstance(dest_min, np.ndarray):
        dest_min_temp = dest_min[band_i]
    else:
        dest_min_temp = dest_min

    if isinstance(dest_max, list) or isinstance(dest_max, np.ndarray):
        dest_max_temp = dest_max[band_i]
    else:
        dest_max_temp = dest_max

    return (src_min_temp, src_max_temp, dest_min_temp, dest_max_temp)
